import os
from utils.logger_config import LoggerSingleton

logger = LoggerSingleton.get_logger()

def save_graph_image(graph, base_dir="output/graph_visual"):
    """
    生成 Mermaid 流程图并保存到指定目录，文件名为 "graph_{id}.png"，
    其中 id 为当前目录中已有图片的数量。

    参数：
    - graph: 已编译的 LangGraph 图对象。
    - base_dir: 图片保存的目录（默认为 "output/graph_visual"）。

    返回：
    - 保存的文件路径。
    """
    try:
        # 确保 base_dir 目录存在
        os.makedirs(base_dir, exist_ok=True)

        # 获取 base_dir 中的所有 .png 文件，并按文件数量生成新的 id
        existing_files = [f for f in os.listdir(base_dir) if f.endswith(".png")]
        id = len(existing_files) + 1  # 新的文件 id 是当前目录中已有文件的数量 + 1

        # 生成新的文件名
        filename = f"graph_{id}.png"
        save_path = os.path.join(base_dir, filename)

        # 生成 Mermaid PNG 图片
        img_data = graph.get_graph().draw_mermaid_png()

        # 保存图片
        with open(save_path, "wb") as f:
            f.write(img_data)

        logger.info(f"流程图已成功保存到: {save_path}")
        return save_path
    except Exception as e:
        logger.error(f"保存流程图失败: {e}")
        return None
